package ssi.staaar.ciphers;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.Serializable;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.EncodedKeySpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.Cipher;
import javax.crypto.CipherOutputStream;
import javax.crypto.NoSuchPaddingException;

import ssi.peeno.model.User;
import encoders.Base64;

public class AsymCipher implements Serializable {
	private RSAPublicKey publicKey;
	private RSAPrivateKey privateKey;
	private String userName;
	private Cipher cipher;
	private String pwd;
	{
		init();
	}
	
	public AsymCipher(String user) {
		this.userName = user;
	}
	
	public AsymCipher(User user) {
		publicKey = (RSAPublicKey)user.getPubKey();
	}
	
	public AsymCipher(String user, String pwd) {
		this.userName = user;
		this.pwd = pwd;
	}
	
	public AsymCipher() {
	}

	private void init() {
		try {
			cipher = Cipher.getInstance("RSA/None/NoPadding", "BC");
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (NoSuchProviderException e) {
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			e.printStackTrace();
		}
	}
	
	/**
	 * Crypt the plaintext with key. 
	 * @param plaintext
	 * @param key
	 * @return String
	 * @throws Exception
	 */
	public String crypt(String plaintext, Key key) throws Exception {
		byte[] cipherTextBytes = null;
		cipher.init(Cipher.ENCRYPT_MODE, key);
		byte[] plaintextBytes=plaintext.getBytes("UTF-8");
		//split in block and crypt
		cipherTextBytes=splitAndDoIt(cipher, plaintextBytes);	
		String cipherText = new String(Base64.encode((cipherTextBytes)));
		return cipherText;
	}
	
	public String cryptWithPublicKey(String plaintext) throws Exception {
		return crypt(plaintext, publicKey);
	}
	
	public String cryptWithPrivateKey(String plaintext) throws Exception {
		return crypt(plaintext, privateKey);
	}
	
	public String decrypt(String ciphertext, Key key) throws Exception {
		byte[] cipherTextBytes = Base64.decode(ciphertext);
		byte[] decryptedTextBytes = null;
		cipher.init(Cipher.DECRYPT_MODE, key); 
		//split in block and decrypt
		decryptedTextBytes=splitAndDoIt(cipher, cipherTextBytes);
		String decryptedText = new String(decryptedTextBytes, "UTF-8");
		return decryptedText;
	}

	public String decryptWithPrivate(String ciphertext) throws Exception {
		return decrypt(ciphertext, privateKey);
	}
	
	public String decryptWithPublic(String ciphertext)throws Exception {
		return decrypt(ciphertext, publicKey);
	}
	
	private byte[] splitAndDoIt(Cipher cipher, byte[] blocks)
			throws IOException {
		byte[] textBytes;
		ByteArrayOutputStream cipherTextBaos = new ByteArrayOutputStream();
		ByteArrayOutputStream baos;
		CipherOutputStream cos;
		int cipherTextLength = blocks.length;
		int blockSize = cipher.getBlockSize(); // get size of block I have to
		// split
		for (int bytesRead = 0; bytesRead < cipherTextLength; bytesRead = bytesRead
				+ blockSize) {
			if (bytesRead + blockSize > cipherTextLength) {
				// process and write in baos the processed byte LAST BLOCK
				baos = new ByteArrayOutputStream();
				cos = new CipherOutputStream(baos, cipher);
				cos.write(blocks, bytesRead, cipherTextLength - bytesRead);
				cos.close();
				cipherTextBaos.write(baos.toByteArray());
			} else {
				// process and write in baos the processed byte
				baos = new ByteArrayOutputStream();
				cos = new CipherOutputStream(baos, cipher);
				cos.write(blocks, bytesRead, blockSize);
				cos.close();
				cipherTextBaos.write(baos.toByteArray());
			}
		}
		cipherTextBaos.close();
		textBytes = cipherTextBaos.toByteArray();
		return textBytes;
	}

	/**
	 * Generate RSA public and private keys
	 */
	public void generateKey() throws NoSuchAlgorithmException, NoSuchProviderException {
		KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA", "BC");
		kpg.initialize(1024);
		KeyPair keys = kpg.generateKeyPair();
		publicKey = (RSAPublicKey) keys.getPublic();
		privateKey = (RSAPrivateKey) keys.getPrivate();
	}

	public Key getPrivateKey() {
		return privateKey;
	}

	public Key getPublicKey() {
		return publicKey;
	}
	
	public void setPublicKey(Key publicKey) {
		this.publicKey = (RSAPublicKey) publicKey;
	}
	
	public void setPrivateKey(Key privateKey) {
		this.privateKey = (RSAPrivateKey) privateKey;
	}
	
	/**
	 * Get public key from string encoded form of public key
	 * @param keyString
	 * @return PublicKey
	 * @throws Exception
	 */
	public static PublicKey getPubKeyFromString(String keyString) throws Exception {		
		System.out.println(keyString.length());
		byte[] decodedKey = Base64.decode(keyString);
		KeyFactory keyFactory = KeyFactory.getInstance("RSA", "BC");
		EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(decodedKey);
		PublicKey publicKeyDecoded = (PublicKey) keyFactory.generatePublic(publicKeySpec);
		return publicKeyDecoded;
	}
	
	/**
	 * Get private key from string encoded form of public key
	 * @param keyString
	 * @return PrivateKey
	 * @throws Exception
	 */
	public static PrivateKey getPrivKeyFromString(String priv) throws Exception {
		byte[] decodedKey = Base64.decode(priv);
		PrivateKey decoded = null;
		KeyFactory keyFactory;
		try {
			keyFactory = KeyFactory.getInstance("RSA", "BC");
			EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(decodedKey);
			decoded = (PrivateKey) keyFactory.generatePrivate(privateKeySpec);
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (NoSuchProviderException e) {
			e.printStackTrace();
		} catch (InvalidKeySpecException e) {
			e.printStackTrace();
		}
		return decoded;
	}
}
